{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Afl-y-M6k1sk"
   },
   "source": [
    "# Tutorial of merging camera data and extracting camera features for WOMD\n",
    "\n",
    "This tutorial demonstrates how to add camera tokens to the original WOMD scenes. It also provides methods to extract camera features from the merged scenario proto message. Note that WOMD also provides APIs to load the LiDAR data in the tutorial `tutorial_womd_lidar.ipynb`.\n",
    "\n",
    "## Install\n",
    "\n",
    "To run Jupyter Notebook locally:\n",
    "\n",
    "```\n",
    "python3 -m pip install waymo-open-dataset-tf-2-12-0==1.6.7\n",
    "python3 -m pip install \"notebook>=5.3\" \"ipywidgets>=7.5\"\n",
    "python3 -m pip install --upgrade jupyter_http_over_ws>=0.0.7 && \\\n",
    "jupyter serverextension enable --py jupyter_http_over_ws\n",
    "jupyter notebook\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "03uEfb7cCodZ"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "tf.enable_eager_execution()\n",
    "\n",
    "from waymo_open_dataset import dataset_pb2\n",
    "from waymo_open_dataset.protos import scenario_pb2\n",
    "from waymo_open_dataset.utils import womd_camera_utils"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Yt8jNg3IC3Bu"
   },
   "source": [
    "# Augmenting a WOMD scenario\n",
    "\n",
    "To augment the original WOMD with camera data for input frames, there are three steps:\n",
    "1. Load the first scenario proto message from the motion data and check the `scenario_id` field.\n",
    "2. Find the corresponding frame camera data file which has the file name as `{scenario_id}.tfrecord`.\n",
    "3. Load the frame camera data file which is a scenario proto with non-empty `frame_camera_tokens` field only and merge the loaded data into the scenario proto's `frame_camera_tokens` field."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "6uCajVOzGFaC"
   },
   "outputs": [],
   "source": [
    "def _load_scenario_data(tfrecord_file: str) -> scenario_pb2.Scenario:\n",
    "  \"\"\"Load a scenario proto from a tfrecord dataset file.\"\"\"\n",
    "  dataset = tf.data.TFRecordDataset(tfrecord_file, compression_type='')\n",
    "  data = next(iter(dataset))\n",
    "  return scenario_pb2.Scenario.FromString(data.numpy())\n",
    "\n",
    "WOMD_FILE = '/content/waymo-od/tutorial/womd_scenario_input.tfrecord'\n",
    "womd_original_scenario = _load_scenario_data(WOMD_FILE)\n",
    "print(f'Loaded a scenario with the scenario_id {womd_original_scenario.scenario_id}'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "tfTb3ZoaGLCv"
   },
   "outputs": [],
   "source": [
    "# The corresponding compressed camera data file has the name\n",
    "# {scenario_id}.tfrecord. For simplicity, we rename the corresponding camera\n",
    "# data file 'ee519cf571686d19.tfrecord' to be\n",
    "# 'womd_lidar_and_camera_data.tfrecord'.\n",
    "CAMERA_DATA_FILE = '/content/waymo-od/tutorial/womd_lidar_and_camera_data.tfrecord'\n",
    "womd_camera_scenario = _load_scenario_data(CAMERA_DATA_FILE)\n",
    "scenario_augmented = womd_camera_utils.add_camera_tokens_to_scenario(\n",
    "    womd_original_scenario, womd_camera_scenario)\n",
    "print(f'#frames = {len(scenario_augmented.frame_camera_tokens)}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "uzhDpVHFEEst"
   },
   "source": [
    "# Extract camera features\n",
    "\n",
    "The camera data in the WOMD proto files is a sequence of integers for each frame and each sensor. Each sensor image is encoded with 256 integers (tokens). The integer refers to the row index of a pre-trained coodebook. We provide the codebook in the tutorial and show how to extract the corresponding camera features."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "3DXErUfxsCiE"
   },
   "outputs": [],
   "source": [
    "WOMD_CAMERA_CODEBOOK_FILE = '/content/waymo-od/tutorial/womd_camera_codebook.npy'\n",
    "womd_camera_codebook = np.load(WOMD_CAMERA_CODEBOOK_FILE)\n",
    "\n",
    "cur_frame_index = 0\n",
    "for camera_tokens in scenario_augmented.frame_camera_tokens[cur_frame_index].camera_tokens:\n",
    "  print(f'Camera name = {camera_tokens.camera_name}')\n",
    "  tokens = np.array(camera_tokens.tokens, dtype=int)\n",
    "  embedding = womd_camera_utils.get_camera_embedding_from_codebook(\n",
    "      womd_camera_codebook, tokens\n",
    "  )\n",
    "  print(f'Embedding shape = {embedding.shape}')"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "private_outputs": true
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
